#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "types.h"
#include "state_vector.h"
#include "convective.h"

params_t *params;

//------------------------------------------------------------------------------

void save_bin(sv_t *sv) {

  FILE *fp;
  char file_name[128];
  static int count;

  sprintf(file_name, "output_%06i.bin", count++);

  fp = fopen(file_name, "wb");

  fwrite(params, sizeof(params_t), 1, fp);

  sv_fwrite(sv, fp);

  fclose(fp);

}

//------------------------------------------------------------------------------

void init() {

  params = calloc(1, sizeof(params_t));

  params->Re = 1.0;
  params->Nr = 10;
  params->Nt = 10;
  params->Nz = 10;
  params->Lr = 1.0;
  params->Lz = 1.0;
  params->dt = 1.0e-4;
  params->dr = params->Lr / params->Nr;
  params->dz = params->Lz / params->Nz;
  params->dtheta = 2.0 * PI / params->Nt;

}

//------------------------------------------------------------------------------

void build_domain(sv_t *sv) {

  int i, j, k, s;

  for ( k = 0 ; k < params->Nz ; k++ ) {
	for ( j = 0 ; j < params->Nt ; j++ ) {
	  for ( i = 0 ; i < params->Nr ; i++ ) {

		s = i + params->Nr * ( j + params->Nt * k );

		if ( i < params->Nr / 10 ) {

		  sv->type[s] = CELL_BOUNDARY;

		} else {

		  sv->type[s] = CELL_FULL;

		  sv->vt[s] = 1.0;
		}

	  }
	}
  }

}

//------------------------------------------------------------------------------

void save_vtk() {

  FILE *fp_in, *fp_out;
  sv_t *sv;
  int i, j, k, s;
  int count;
  double dr, dteta, dz;
  double x, y, z, r, teta;
  double vx, vy, vz, vr, vt;
  char file_name_in[128], file_name_out[128];

  params = calloc(1, sizeof(params_t));

  for ( count = 0 ; count < 1000000 ; count++ ) {

	sprintf(file_name_in, "output_%06i.bin", count);
	sprintf(file_name_out, "output_%06i.vtk", count);

	if ( ( fp_in = fopen(file_name_in, "rb") ) == 0 ) {
	  break;
	}

	fp_out = fopen(file_name_out, "wb");

	if ( fread(params, sizeof(params_t), 1, fp_in) != 1 ) {

	  printf("save_vtk error\n");
	  exit(0);

	}

	sv = sv_alloc();

	sv_fread(sv, fp_in);

	fprintf(fp_out, "# vtk DataFile Version 2.0\n");
	fprintf(fp_out, "Unstructured Grid Example\n");
	fprintf(fp_out, "ASCII\n");
	fprintf(fp_out, "DATASET UNSTRUCTURED_GRID\n");
	fprintf(fp_out, "POINTS %i double\n", (params->Nr+1) * params->Nt * (params->Nz+1));

	dr    = params->Lr / params->Nr;
	dteta = 2.0 * PI   / params->Nt;
	dz    = params->Lz / params->Nz;

	for ( k = 0 ; k <= params->Nz ; k++ ) {

	  z = k * dz;

	  for ( j = 0 ; j < params->Nt ; j++ ) {

		teta = j * dteta;

		for ( i = 0 ; i <= params->Nr ; i++ ) {

		  r = i * dr;

		  x = r * cos(teta);
		  y = r * sin(teta);

		  fprintf(fp_out, "%e %e %e\n", x, y, z);

		}
	  }
	}

	fprintf(fp_out, "CELLS %i %i\n", params->Nr * params->Nt * params->Nz, 9 * params->Nr * params->Nt * params->Nz);

	for ( k = 0 ; k < params->Nz ; k++ ) {
	  for ( j = 0 ; j < params->Nt ; j++ ) {
		for ( i = 0 ; i < params->Nr ; i++ ) {

		  fprintf(fp_out, "8 ");
		  fprintf(fp_out, "%i ", (i+0) + (params->Nr+1) * ( (j+0) % params->Nt + params->Nt * (k+0) ) );
		  fprintf(fp_out, "%i ", (i+1) + (params->Nr+1) * ( (j+0) % params->Nt + params->Nt * (k+0) ) );
		  fprintf(fp_out, "%i ", (i+1) + (params->Nr+1) * ( (j+1) % params->Nt + params->Nt * (k+0) ) );
		  fprintf(fp_out, "%i ", (i+0) + (params->Nr+1) * ( (j+1) % params->Nt + params->Nt * (k+0) ) );
		  fprintf(fp_out, "%i ", (i+0) + (params->Nr+1) * ( (j+0) % params->Nt + params->Nt * (k+1) ) );
		  fprintf(fp_out, "%i ", (i+1) + (params->Nr+1) * ( (j+0) % params->Nt + params->Nt * (k+1) ) );
		  fprintf(fp_out, "%i ", (i+1) + (params->Nr+1) * ( (j+1) % params->Nt + params->Nt * (k+1) ) );
		  fprintf(fp_out, "%i ", (i+0) + (params->Nr+1) * ( (j+1) % params->Nt + params->Nt * (k+1) ) );
		  fprintf(fp_out, "\n");

		}
	  }
	}

	fprintf(fp_out, "CELL_TYPES %i\n", params->Nr * params->Nt * params->Nz);

	for ( k = 0 ; k < params->Nz ; k++ ) {
	  for ( j = 0 ; j < params->Nt ; j++ ) {
		for ( i = 0 ; i < params->Nr ; i++ ) {

		  fprintf(fp_out, "12\n");

		}
	  }
	}

	fprintf(fp_out, "CELL_DATA %i\n", params->Nr * params->Nt * params->Nz);
	fprintf(fp_out, "SCALARS cell_type int 1\n");
	fprintf(fp_out, "LOOKUP_TABLE default\n");

	for ( k = 0 ; k < params->Nz ; k++ ) {
	  for ( j = 0 ; j < params->Nt ; j++ ) {
		for ( i = 0 ; i < params->Nr ; i++ ) {

		  s = i + params->Nr * ( j + params->Nt * k );

		  fprintf(fp_out, "%i\n", sv->type[s]);

		}
	  }
	}

	fprintf(fp_out, "VECTORS velocity double\n");

	for ( k = 0 ; k < params->Nz ; k++ ) {
	  for ( j = 0 ; j < params->Nt ; j++ ) {
		for ( i = 0 ; i < params->Nr ; i++ ) {

		  s = (i+0) + params->Nr * ( (j+0) % params->Nt + params->Nt * (k+0) );
		  vr = sv->vr[s];
		  s = (i+1) + params->Nr * ( (j+0) % params->Nt + params->Nt * (k+0) );
		  vr = 0.5 * ( vr + sv->vr[s] );

		  s = (i+0) + params->Nr * ( (j+0) % params->Nt + params->Nt * (k+0) );
		  vt = sv->vt[s];
		  s = (i+0) + params->Nr * ( (j+1) % params->Nt + params->Nt * (k+0) );
		  vt = 0.5 * ( vt + sv->vt[s] );

		  s = (i+0) + params->Nr * ( (j+0) % params->Nt + params->Nt * (k+0) );
		  vz = sv->vz[s];
		  s = (i+0) + params->Nr * ( (j+0) % params->Nt + params->Nt * (k+1) );
		  vz = 0.5 * ( vz + sv->vz[s] );

		  teta = (j+0.5) * dteta;

		  vx = vr * cos(teta) - vt * sin(teta);
		  vy = vr * sin(teta) + vt * cos(teta);

		  fprintf(fp_out, "%e %e %e\n", vx, vy, vz);

		}
	  }
	}

	fclose(fp_in);
	fclose(fp_out);

	sv_free(sv);

  }

}

//------------------------------------------------------------------------------

double viscous_r(sv_t *y, int i, int j, int k) {

  double temp;

  temp = 1.0 / calc_r(i) / params->dr / params->dr * ( calc_r(i+0.5) * sv_get_vr(y, i+1, j, k) - 2.0 * calc_r(i) * sv_get_vr(y, i, j, k) + calc_r(i-0.5) * sv_get_vr(y, i-1, j, k) );

  temp += 1.0 / calc_r(i) / calc_r(i) / params->dtheta / params->dtheta * ( sv_get_vr(y, i, j+1, k) - 2.0 * sv_get_vr(y, i, j, k) + sv_get_vr(y, i, j-1, k) );

  temp += 1.0 / params->dz / params->dz * ( sv_get_vr(y, i, j, k+1) - 2.0 * sv_get_vr(y, i, j, k) + sv_get_vr(y, i, j, k-1) );

  temp -= 1.0 / calc_r(i) / calc_r(i) * sv_get_vr(y, i, j, k);

  temp -= 1.0 / calc_r(i) / calc_r(i) / params->dtheta * ( sv_get_vt(y, i, j+1, k) + sv_get_vt(y, i-1, j+1, k) - sv_get_vt(y, i, j, k) - sv_get_vt(y, i-1, j, k) );

  return temp / params->Re;

}

//------------------------------------------------------------------------------

double viscous_t(sv_t *y, int i, int j, int k) {

  double temp;

  temp = 1.0 / calc_r(i+0.5) / params->dr / params->dr * ( calc_r(i+1) * sv_get_vt(y, i+1, j, k) - 2.0 * calc_r(i+0.5) * sv_get_vt(y, i, j, k) + calc_r(i) * sv_get_vt(y, i-1, j, k) );

  temp += 1.0 / calc_r(i+0.5) / calc_r(i+0.5) / params->dtheta / params->dtheta * ( sv_get_vt(y, i, j+1, k) - 2.0 * sv_get_vt(y, i, j, k) + sv_get_vt(y, i, j-1, k) );

  temp += 1.0 / params->dz / params->dz * ( sv_get_vt(y, i, j, k+1) - 2.0 * sv_get_vt(y, i, j, k) + sv_get_vt(y, i, j, k-1) );

  temp -= 1.0 / calc_r(i+0.5) / calc_r(i+0.5) * sv_get_vt(y, i, j, k);

  temp += 1.0 / calc_r(i+0.5) / calc_r(i+0.5) / params->dtheta * ( sv_get_vr(y, i+1, j, k) + sv_get_vr(y, i, j, k) - sv_get_vr(y, i+1, j-1, k) - sv_get_vr(y, i, j-1, k) );

  return temp / params->Re;

}

//------------------------------------------------------------------------------

double viscous_z(sv_t *y, int i, int j, int k) {

  double temp;

  temp = 1.0 / calc_r(i+0.5) / params->dr / params->dr * ( calc_r(i+1) * sv_get_vz(y, i+1, j, k) - 2.0 * calc_r(i+0.5) * sv_get_vz(y, i, j, k) + calc_r(i) * sv_get_vz(y, i-1, j, k) );

  temp += 1.0 / calc_r(i+0.5) / calc_r(i+0.5) / params->dtheta / params->dtheta * ( sv_get_vz(y, i, j+1, k) - 2.0 * sv_get_vz(y, i, j, k) + sv_get_vz(y, i, j-1, k) );

  temp += 1.0 / params->dz / params->dz * ( sv_get_vz(y, i, j, k+1) - 2.0 * sv_get_vz(y, i, j, k) + sv_get_vz(y, i, j, k-1) );

  return temp / params->Re;

}

//------------------------------------------------------------------------------

// Momentum Equation

sv_t *calc_f(double t, sv_t *y) {

  int i, j, k;
  sv_t *z = sv_dup(y);

  t = t;

#pragma omp parallel for private(i, j, k)

  for ( j = 0 ; j < params->Nt ; j++ ) {

	for ( k = 0 ; k < params->Nz ; k++ ) {
	  for ( i = 0 ; i <= params->Nr ; i++ ) {

		if ( sv_get_type(y, i, j, k) == CELL_FULL && sv_get_type(y, i-1, j, k) == CELL_FULL ) {

		  sv_set_vr(z, i, j, k, viscous_r(y, i, j, k) + convective_central_r(y, i, j, k) );


		}

	  }
	}

	for ( k = 0 ; k < params->Nz ; k++ ) {
	  for ( i = 0 ; i < params->Nr ; i++ ) {

		if ( sv_get_type(y, i, j, k) == CELL_FULL && sv_get_type(y, i, j-1, k) == CELL_FULL ) {

		  sv_set_vt(z, i, j, k, viscous_t(y, i, j, k) + convective_central_t(y, i, j, k) );

		}

	  }
	}

	for ( k = 0 ; k <= params->Nz ; k++ ) {
	  for ( i = 0 ; i < params->Nr ; i++ ) {

		if ( sv_get_type(y, i, j, k) == CELL_FULL && sv_get_type(y, i, j, k-1) == CELL_FULL ) {

		  sv_set_vz(z, i, j, k, viscous_z(y, i, j, k) + convective_central_z(y, i, j, k) );

		}

	  }
	}

  }

  return z;

}

//------------------------------------------------------------------------------

// Explicit 4th Order Runge Kutta Method

sv_t *rk4(sv_t *y) {

  sv_t *k1, *k2, *k3, *k4;
  sv_t *temp = sv_dup(y);
  double t = params->t;
  double h = params->dt;

  k1 = calc_f(t, y);

  sv_axpby(temp, 1.0, y, 0.5*h, k1);
  k2 = calc_f(t+0.5*h, temp);

  sv_axpby(temp, 1.0, y, 0.5*h, k2);
  k3 = calc_f(t+0.5*h, temp);

  sv_axpby(temp, 1.0, y, h, k3);
  k4 = calc_f(t+h, temp);

  sv_axpby(temp, 1.0,   k1,   2.0, k2);
  sv_axpby(temp, 1.0,   temp, 2.0, k3);
  sv_axpby(temp, 1.0,   temp, 1.0, k4);
  sv_axpby(temp, h/6.0, temp, 1.0, y);

  sv_free(k1);
  sv_free(k2);
  sv_free(k3);
  sv_free(k4);

  return temp;

}

//------------------------------------------------------------------------------

void solve(sv_t *y0) {

  sv_t *y;

  do {

	save_bin(y0);

	y = rk4(y0);

	params->t += params->dt;

	(params->cycle)++;

	sv_free(y0);

	y0 = y;

  } while ( params->cycle < 100 );

}

//------------------------------------------------------------------------------

int main(int argc, char **argv) {

  sv_t *sv;

  if ( argc == 2 ) {

	if ( strcmp("simulate", argv[1]) == 0 ) {

	  printf("solving...\n");

	  init();

	  sv = sv_alloc();

	  build_domain(sv);

	  solve(sv);

	  sv_free(sv);

	  return 0;

	} else if ( strcmp("vtk", argv[1]) == 0 ){

	  printf("saving vtk files...\n");

	  save_vtk();

	  return 0;

	}

  }

  printf("invalid sintax\n");

  return 0;

}

//------------------------------------------------------------------------------

